package crawl

import (
	"io"
	"net/http"
	"os"
	"time"
)

func ProcessProgress(readBody func(body io.Reader) error, report func(cur, total int64)) Processor {
	return func(resp *http.Response) (err error) {
		pw := &progressWriter{total: resp.ContentLength, report: report}
		defer pw.update()
		return readBody(io.TeeReader(resp.Body, pw))
	}
}

func ProcessDownload(fn string, report func(cur, total int64)) Processor {
	return ProcessProgress(func(body io.Reader) (err error) {
		var dst io.WriteCloser
		if dst, err = os.Create(fn); err != nil {
			return err
		}
		if _, err = io.Copy(dst, body); err != nil {
			_ = os.Remove(fn)
		}
		return
	}, report)
}

type progressWriter struct {
	w        io.Writer
	cur      int64
	total    int64
	report   func(cur, total int64)
	reportAt int64
}

func (pw *progressWriter) Write(p []byte) (n int, err error) {
	if pw.w == nil {
		n = len(p)
	} else if n, err = pw.w.Write(p); err != nil {
		return
	}

	pw.cur += int64(n)
	if now := time.Now().Unix(); now > pw.reportAt || pw.cur == pw.total {
		pw.update()
		pw.reportAt = now
	}
	return
}

func (pw *progressWriter) update() {
	if pw.report != nil {
		pw.report(pw.cur, pw.total)
	}
}
